package com.atguigu.ai.operator;

import cn.hutool.core.util.IdUtil;
import com.atguigu.ai.annotation.BeanType;
import com.atguigu.ai.utils.MongoUtil;
import com.atguigu.ai.vo.QueryVo;
import com.atguigu.ai.enums.AiTypeEnum;
import com.atguigu.ai.enums.LanguageEnum;
import com.atguigu.ai.enums.SystemConstant;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import reactor.core.publisher.Flux;

import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@BeanType(AiTypeEnum.OPENAI)
public class OpenAiOperator implements AiOperator {

    @Autowired
    private OpenAiChatModel openAiChatModel;

    @Autowired
    // private RedisVectorStore openaiRedisVectorStore;
    private QdrantVectorStore openaiQdrantVectorStore;

    @Autowired
    private MongoTemplate mongoTemplate;

    @Override
    public String chat(QueryVo queryVo) {
        return null;
    }

    @Override
    public Flux<String> chatStream(String model, QueryVo queryVo) {
        // 把问题记录到mongodb
        com.atguigu.ai.pojo.Message msg = new com.atguigu.ai.pojo.Message();
        msg.setChatId(queryVo.getChatId());
        msg.setType(0);
        msg.setContent(queryVo.getMsg());
        msg.setCreateTime(new Date());
        msg.setId(IdUtil.getSnowflake().nextId());
        this.mongoTemplate.insert(msg, MongoUtil.getMessageCollection(queryVo.getChatId()));

        // 查询本地知识库
        List<Document> results = openaiQdrantVectorStore.similaritySearch(SearchRequest
                .query(queryVo.getMsg())
                .withFilterExpression(
                        new FilterExpressionBuilder()
                                .eq("projectId", queryVo.getProjectId()) // 查询当前用户及管理员的本地知识库
                                .build())
                .withTopK(SystemConstant.TOPK)); // 取前10个

        // 把本地知识库的内容作为系统提示放入
        List<Message> msgList = results.stream().map(result ->
                new SystemMessage(result.getContent())).collect(Collectors.toList());
        // 中英文切换
        msgList.add(new SystemMessage(LanguageEnum.getMsg(queryVo.getLanguage())));
        // 加入当前用户的提问
        msgList.add(new UserMessage(queryVo.getMsg()));

        // 提交到大模型获取最终结果
        Flux<ChatResponse> responseFlux = this.openAiChatModel.stream(
                new Prompt(msgList, OpenAiChatOptions.builder().withModel(model).build()));
        Flux<String> flux = responseFlux.map(response -> response.getResult() != null
                && response.getResult().getOutput() != null
                && response.getResult().getOutput().getContent() != null
                ? response.getResult().getOutput().getContent() : "");

        // flux.collectList().subscribe(list -> {\
        //     此处获取的信息和最终返回的信息 是两个结果
        //     System.out.println(StringUtils.join(list, ""));
        // });
        return flux;
    }

    @Override
    public String imageUrl(QueryVo queryVo) {
        return null;
    }

    @Override
    public String imageBase64Json(QueryVo queryVo) {
        return null;
    }

    @Override
    public String textToAudio(QueryVo queryVo) {
        return null;
    }

    @Override
    public Boolean upload(Long projectId, Long knowledgeId, String content) {
        Document document = new Document(knowledgeId.toString(), content, Map.of("projectId", projectId));
        this.openaiQdrantVectorStore.add(List.of(document));
        return true;
    }

    @Override
    public Boolean remove(String docId) {
        return this.openaiQdrantVectorStore.delete(List.of(docId)).get();
    }

    @Override
    public Flux<String> stream(QueryVo queryVo) {
        return this.openAiChatModel.stream(queryVo.getMsg());
    }
}
